import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough, RunnableLambda
from langchain_core.output_parsers import StrOutputParser

model = ChatOpenAI(
    model="deepseek-chat",
    openai_api_key=os.getenv("DEEPSEEK_API_KEY"),
    openai_api_base=os.getenv("DEEPSEEK_API_BASE")
)

#使用提示词模板
# --- 新增代码：使用 RunnableParallel 根据标题生成文案 ---
print("\n--- RunnableParallel 案例：根据标题生成文案 ---")

# 1. 定义文章生成的提示模板
prompt = ChatPromptTemplate.from_template(
    "请为以下标题写一篇简短的文章（大约100-200字）：\n\n标题：{title}\n\n文章内容："
)

# 2. 定义文章生成链
# 这个链接收一个字典 {"title": "..."} 作为输入 (因为 ChatPromptTemplate 的 input_variables 是 {"title"})
# 或者如果直接用 .invoke("some title") 调用 prompt_for_article, LangChain 会自动包装成 {"title": "some title"}
# 但在这里，我们将把它作为 RunnableParallel 的一部分，它会正确处理输入。
chain = prompt | model | StrOutputParser()

# 3. 使用 RunnableParallel 组合标题和生成的文章内容
# 当 copy_creation_chain 被以一个标题字符串（例如 "人工智能的未来"）调用时：
# - article_generation_chain 会接收这个标题字符串作为输入，并生成文章内容。
# - RunnablePassthrough() 会直接传递这个标题字符串。
combine_elements_step = RunnableParallel(
    article_content=chain,  # 输入是标题字符串，输出是文章内容字符串
    original_title=RunnablePassthrough()       # 输入是标题字符串，输出是标题字符串
)

# 4. 定义最终文案的格式化函数
def format_final_copy(data: dict) -> str:
    """将标题和文章内容合并为最终文案。"""
    return f"标题：{data['original_title']}\n\n文章内容：\n{data['article_content']}"

# 将格式化函数包装成 RunnableLambda
final_formatter = RunnableLambda(format_final_copy)

# 5. 构建完整的文案生成链
# 链的流程:
# Input: "一个标题字符串"
# 1. combine_elements_step (input: "一个标题字符串")
#    - article_generation_chain (input: "一个标题字符串") -> "生成的文章内容..."
#    - RunnablePassthrough (input: "一个标题字符串")    -> "一个标题字符串"
#    Output: {"article_content": "生成的文章内容...", "original_title": "一个标题字符串"}
# 2. final_formatter (input: 字典) -> "格式化后的完整文案"
copy_creation_chain = combine_elements_step | final_formatter

# 6. 调用链并打印结果
example_title = "LangChain中RunnableParallel的妙用"
print(f"正在为标题生成文案：'{example_title}'")

final_copy = copy_creation_chain.invoke(example_title)
print("\n--- 生成的完整文案 ---")
print(final_copy)

example_title_2 = "探索宇宙的奥秘"
print(f"\n正在为标题生成文案：'{example_title_2}'")
final_copy_2 = copy_creation_chain.invoke(example_title_2)
print("\n--- 生成的完整文案 ---")
print(final_copy_2)




